# Single Shot MultiBox Detector
import mxnet.ndarray as nd
import mxnet.image as image

import matplotlib as mpl
import matplotlib.pyplot as plt

mpl.rcParams['figure.dpi'] = 120
data_shape = 256
batch_size = 32
rgb_mean = nd.array([123, 117, 104])
DATA_DIR = '/home/alpha/ML/data/pikachu/'


def get_iterators(data_dir, data_shape, batch_size):
    class_names = ['pikachu']
    num_class = len(class_names)
    train_iter = image.ImageDetIter(
        batch_size=batch_size,
        data_shape=(3, data_shape, data_shape),
        path_imgrec=data_dir + 'train.rec',
        path_imgidx=data_dir + 'train.idx',
        shuffle=True,
        mean=True,
        rand_crop=1,
        min_object_covered=0.95,
        max_attempts=200)
    val_iter = image.ImageDetIter(
        batch_size=batch_size,
        data_shape=(3, data_shape, data_shape),
        path_imgrec=data_dir + 'val.rec',
        shuffle=False,
        mean=True)
    return train_iter, val_iter, class_names, num_class


def box_to_rect(box, color, linewidth=3):
    """convert an anchor box to a matplotlib rectangle"""
    box = box.asnumpy()
    return plt.Rectangle(
        (box[0], box[1]), box[2] - box[0], box[3] - box[1],
        fill=False, edgecolor=color, linewidth=linewidth)


def test_plot_box():
    train_data, test_data, class_names, num_class = get_iterators(
        DATA_DIR, data_shape, batch_size)

    batch = train_data.next()
    print(batch)
    _, figs = plt.subplots(3, 3, figsize=(6, 6))
    for i in range(3):
        for j in range(3):
            img, labels = batch.data[0][3 * i + j], batch.label[0][3 * i + j]
            # (3L, 256L, 256L) => (256L, 256L, 3L)
            img = img.transpose((1, 2, 0)) + rgb_mean
            img = img.clip(0, 255).asnumpy() / 255
            fig = figs[i][j]
            fig.imshow(img)
            for label in labels:
                rect = box_to_rect(label[1:5] * data_shape, 'red', 2)
                fig.add_patch(rect)
            fig.axes.get_xaxis().set_visible(False)
            fig.axes.get_yaxis().set_visible(False)
    plt.show()


if __name__ == '__main__':
    # train_data, test_data, class_names, num_class = get_iterators(
    #     DATA_DIR, data_shape, batch_size)

    test_plot_box()
    # print(train_data.next())
